#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
iptables 管理服务
"""

import logging
import subprocess
import re
from datetime import datetime

logger = logging.getLogger(__name__)

class IptablesManager:
    """iptables 管理服务"""

    # 代理端口
    PROXY_PORT = 5555

    # 规则注释前缀，用于识别由本系统添加的规则
    RULE_COMMENT_PREFIX = "KubeRelayer-"

    @staticmethod
    def setup_intercept(service_ip, service_port, comment=None, direction='both'):
        """
        设置拦截规则

        Args:
            service_ip: 服务IP地址
            service_port: 服务端口
            comment: 规则注释，用于后续识别和删除
            direction: 拦截方向，可选值为 'input'(入站流量), 'output'(出站流量) 或 'both'(同时拦截入站和出站流量)

        Returns:
            bool: 是否成功设置
        """
        if not service_ip or not service_port:
            logger.error("服务IP或端口不能为空")
            return False

        # 生成规则注释
        rule_comment = f"{IptablesManager.RULE_COMMENT_PREFIX}{comment or datetime.now().strftime('%Y%m%d%H%M%S')}"

        try:
            # 根据方向设置规则
            if direction in ['input', 'both']:
                # 设置 PREROUTING 规则，拦截发往目标服务的请求(入站流量)
                subprocess.run([
                    'iptables', '-t', 'nat', '-A', 'PREROUTING',
                    '-p', 'tcp', '-d', service_ip, '--dport', str(service_port),
                    '-j', 'REDIRECT', '--to-port', str(IptablesManager.PROXY_PORT),
                    '-m', 'comment', '--comment', rule_comment
                ], check=True)
                logger.info(f"成功设置入站拦截规则: {service_ip}:{service_port} -> {IptablesManager.PROXY_PORT}")

            if direction in ['output', 'both']:
                # 设置 OUTPUT 规则，拦截本地发出的请求(出站流量)
                subprocess.run([
                    'iptables', '-t', 'nat', '-A', 'OUTPUT',
                    '-p', 'tcp', '-d', service_ip, '--dport', str(service_port),
                    '-j', 'REDIRECT', '--to-port', str(IptablesManager.PROXY_PORT),
                    '-m', 'comment', '--comment', rule_comment
                ], check=True)
                logger.info(f"成功设置出站拦截规则: {service_ip}:{service_port} -> {IptablesManager.PROXY_PORT}")

            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"设置拦截规则失败: {e}")
            return False

    @staticmethod
    def remove_intercept(service_ip, service_port, comment=None):
        """
        移除拦截规则

        Args:
            service_ip: 服务IP地址
            service_port: 服务端口
            comment: 规则注释，用于识别要删除的规则

        Returns:
            bool: 是否成功移除
        """
        if not service_ip or not service_port:
            logger.error("服务IP或端口不能为空")
            return False

        # 生成规则注释匹配模式
        comment_pattern = f"{IptablesManager.RULE_COMMENT_PREFIX}"
        if comment:
            comment_pattern += comment

        try:
            # 获取当前规则
            nat_rules = subprocess.check_output(['iptables', '-t', 'nat', '-L', '-n', '--line-numbers']).decode('utf-8')

            # 查找匹配的 PREROUTING 规则
            prerouting_rules = re.findall(
                r'^(\d+).*REDIRECT.*dpt:' + str(service_port) + r'.*to:' + str(IptablesManager.PROXY_PORT) +
                r'.*comment_match.*' + comment_pattern,
                nat_rules, re.MULTILINE
            )

            # 查找匹配的 OUTPUT 规则
            output_rules = re.findall(
                r'^(\d+).*REDIRECT.*dpt:' + str(service_port) + r'.*to:' + str(IptablesManager.PROXY_PORT) +
                r'.*comment_match.*' + comment_pattern,
                nat_rules, re.MULTILINE
            )

            # 从高到低删除规则（避免删除后行号变化）
            for rule_num in sorted(prerouting_rules, reverse=True):
                subprocess.run(['iptables', '-t', 'nat', '-D', 'PREROUTING', rule_num], check=True)

            for rule_num in sorted(output_rules, reverse=True):
                subprocess.run(['iptables', '-t', 'nat', '-D', 'OUTPUT', rule_num], check=True)

            logger.info(f"成功移除拦截规则: {service_ip}:{service_port}")
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"移除拦截规则失败: {e}")
            return False

    @staticmethod
    def list_intercepts():
        """
        列出所有拦截规则

        Returns:
            list: 拦截规则列表
        """
        try:
            # 获取当前规则
            nat_rules = subprocess.check_output(['iptables', '-t', 'nat', '-L', '-n']).decode('utf-8')

            # 查找所有 REDIRECT 规则
            redirect_rules = re.findall(
                r'REDIRECT.*dpt:(\d+).*to:(\d+).*comment_match "(' + IptablesManager.RULE_COMMENT_PREFIX + r'.*?)"',
                nat_rules
            )

            # 整理规则信息
            intercepts = []
            for rule in redirect_rules:
                service_port, proxy_port, comment = rule
                intercepts.append({
                    "service_port": int(service_port),
                    "proxy_port": int(proxy_port),
                    "comment": comment,
                    "created_at": IptablesManager._extract_timestamp_from_comment(comment)
                })

            return intercepts
        except subprocess.CalledProcessError as e:
            logger.error(f"列出拦截规则失败: {e}")
            return []

    @staticmethod
    def clear_all_intercepts():
        """
        清除所有拦截规则

        Returns:
            bool: 是否成功清除
        """
        try:
            # 获取当前规则
            nat_rules = subprocess.check_output(['iptables', '-t', 'nat', '-L', '-n', '--line-numbers']).decode('utf-8')

            # 查找所有带有特定注释的规则
            prerouting_rules = re.findall(
                r'^(\d+).*REDIRECT.*comment_match "' + IptablesManager.RULE_COMMENT_PREFIX + r'.*?"',
                nat_rules, re.MULTILINE
            )

            output_rules = re.findall(
                r'^(\d+).*REDIRECT.*comment_match "' + IptablesManager.RULE_COMMENT_PREFIX + r'.*?"',
                nat_rules, re.MULTILINE
            )

            # 从高到低删除规则
            for rule_num in sorted(prerouting_rules, reverse=True):
                subprocess.run(['iptables', '-t', 'nat', '-D', 'PREROUTING', rule_num], check=True)

            for rule_num in sorted(output_rules, reverse=True):
                subprocess.run(['iptables', '-t', 'nat', '-D', 'OUTPUT', rule_num], check=True)

            logger.info("成功清除所有拦截规则")
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"清除拦截规则失败: {e}")
            return False

    @staticmethod
    def _extract_timestamp_from_comment(comment):
        """从注释中提取时间戳"""
        try:
            # 注释格式: KubeRelayer-YYYYMMDDHHMMSS
            timestamp_str = comment.replace(IptablesManager.RULE_COMMENT_PREFIX, "")
            return datetime.strptime(timestamp_str, "%Y%m%d%H%M%S")
        except:
            return None

    @staticmethod
    def get_original_destination(socket_fd):
        """
        获取原始目标信息

        Args:
            socket_fd: socket 文件描述符

        Returns:
            tuple: (ip, port) 原始目标IP和端口
        """
        import socket
        import struct

        try:
            # 获取原始目标信息
            SO_ORIGINAL_DST = 80  # 这是 Linux 特定的 socket 选项
            dst = socket.getsockopt(socket_fd, socket.SOL_IP, SO_ORIGINAL_DST, 16)
            port, raw_ip = struct.unpack_from('!2xH4s', dst)
            ip = socket.inet_ntoa(raw_ip)

            return (ip, port)
        except Exception as e:
            logger.error(f"获取原始目标信息失败: {e}")
            return None
